"""
从TXT文档加载数据，向量化后存储到chroma数据库
"""
import os

import dotenv
from langchain_community.document_loaders import TextLoader

from langchain_community.vectorstores import Chroma,Qdrant
#from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

def fun_0():
    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL")
        , base_url=os.getenv("CLOSE_BASE_URL")
        , api_key=os.getenv("CLOSE_API_KEY")
    )

    chroma_db = Chroma(persist_directory="./chroma_db/chroma-1", embedding_function=embedding_model)

    # 加载TXT文档
    loader = TextLoader(file_path="F:/data/1.txt", encoding="utf-8")
    documents = loader.load()
    # print(documents[0].page_content)

    # 创建文本分块器
    test_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    chunks = test_splitter.split_documents(documents)

    try:
        print("================================")
        chroma_db.add_documents(chunks)
        print("向量数据库创建成功")
    except Exception as e:
        print(f"添加文档时发生错误: {e}")
        import traceback
        traceback.print_exc()

fun_0()


def fun_1():
    dotenv.load_dotenv()

    # 加载TXT文档
    loader = TextLoader(file_path="F:/data/1.txt", encoding="utf-8")
    documents = loader.load()
    #print(documents[0].page_content)

    # 创建文本分块器
    test_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    chunks = test_splitter.split_documents(documents)
    # for i , chunk in enumerate(chunks): # enumerate() 函数用于将一个可遍历的数据对象组合为一个索引序列，同时列出数据对象。
    #     print(f"Chunk {i}: {chunk}")


    # 创建嵌入模型
    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL")
        , base_url=os.getenv("CLOSE_BASE_URL")
        , api_key=os.getenv("CLOSE_API_KEY")
    )

    # 创建向量数据库
    # 注意：向量数据库不仅仅存储向量，还存储数据本身
    # pip install -U langchain-chroma==0.1.1 -i https://pypi.tuna.tsinghua.edu.cn/simple/
    chroma_db = Chroma.from_documents(
        documents=chunks
        , embedding=embedding_model
        , persist_directory="./chroma_db/chroma-1"  # 向量数据库的保存路径。如果没有指定路径，则默认存储在内存中。
    )
    print("向量数据库创建成功")
    print(f"向量数据库中存储了 {len(chunks)} 个向量")

#fun_1()

#检索数据库
def fun_2():
    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL")
        , base_url=os.getenv("CLOSE_BASE_URL")
        , api_key=os.getenv("CLOSE_API_KEY")
    )

    chroma_db = Chroma(persist_directory="./chroma_db/chroma-1", embedding_function=embedding_model)
    query = "北京"
    query_vector = embedding_model.embed_query(query)
    print(query_vector)

    results = chroma_db.similarity_search_by_vector(query_vector)
    print(len(results))
    print(f"查询结果数量: {len(results)}")  # 添加数量统计
    for i, result in enumerate(results):
        print(f"结果 {i + 1}: {result.page_content}")  # 打印具体内容


#fun_2()

#检索数据库
def fun_3():
    dotenv.load_dotenv()

    embedding_model = OpenAIEmbeddings(
        model=os.getenv("CLOSE_EMBEDDING_MODEL")
        , base_url=os.getenv("CLOSE_BASE_URL")
        , api_key=os.getenv("CLOSE_API_KEY")
    )

    chroma_db = Chroma(persist_directory="./chroma_db/chroma-1", embedding_function=embedding_model)
    query = "北京"

    results = chroma_db.similarity_search(query)
    print(len(results))
    print(f"查询结果数量: {len(results)}")  # 添加数量统计
    for i, result in enumerate(results):
        print(f"结果 {i + 1}: {result.page_content}")  # 打印具体内容


#fun_3()